# -*- coding: utf-8 -*_
"""
@file name      : model_save.py
@author         : QuZhang
@date           : 2021-1-2 9:49
@brief          : 模型的保存
"""
import torch.nn as nn
from collections import OrderedDict
import torch


class LeNet2(nn.Module):
    def __init__(self, classes):
        super().__init__()

        # 提取特征
        self.features = nn.Sequential(OrderedDict({
            "conv1": nn.Conv2d(3, 6, 5),
            "relu1": nn.ReLU(),
            "maxpool1": nn.MaxPool2d((2, 2), 2),
            "conv2": nn.Conv2d(6, 16, 5),
            "relu2": nn.ReLU(),
            "maxpool2": nn.MaxPool2d(2, 2),
        }))

       # 分类
        self.classifier = nn.Sequential(OrderedDict({
            "linear1": nn.Linear(16*5*5, 120),
            "relu3": nn.ReLU(),
            "linear2": nn.Linear(120, 84),
            "relu4": nn.ReLU(),
            "linear3": nn.Linear(84, classes),
        }))

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size()[0], -1)  # 将一个通道上的数据变成一行,一维张量
        x = self.classifier(x)
        return x

    def initialize(self):
        for p in self.parameters():
            p.data.fill_(20210102)  # 初始化卷积层和全连接层的参数


if __name__ == '__main__':
    net = LeNet2(classes=2021)

    print("训练前：", net.features[0].weight[0, ...])
    net.initialize()
    print("训练后", net.features[0].weight[0, ...])

    # 设置保存路径
    path_model = "./model.pkl"
    path_state_dict = "./model_state_dict.pkl"

    # 保存整个模型
    torch.save(net, path_model)

    # 保存模型参数:一般为卷积层和线性层参数
    net_state_dict = net.state_dict()  # 获取模型参数
    torch.save(net_state_dict, path_state_dict)  # 保存模型参数